package cn.bj.brook.tkit;

import javax.crypto.*;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.SecretKeySpec;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.InputStream;
import java.security.*;
import java.security.cert.CertificateException;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

public class EncDecToolKit {

    /**
     * 此函数为AES加密准备，AES都是以16位字符为一个block进行加密的，所以key和IV（初始向量）都是需要16字符大小
     *
     * @param input 输入的字符串
     * @return 16位字符串
     */
    private static String getPre16OfValue(String input) {
        if (input.length() > 16) {
            return input.substring(0, 16);
        } else {
            int diff = 16 - input.length();
            StringBuilder next = new StringBuilder(input);
            while (diff > 0) {
                next.append("0");
                diff--;
            }
            return next.toString();
        }
    }

    /**
     * AES密码分组链接模式Cipher Block Chaining 加密
     *
     * @param password  输入的对称加密秘钥
     * @param clearText 明文
     * @return 加密的文字base64字符串
     */
    public static String aesCBCEncrypt(String password, String clearText) {
        password = getPre16OfValue(password);
        try {
            Cipher cipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
            SecretKey secKey = new SecretKeySpec(password.getBytes(), "AES");
            cipher.init(Cipher.ENCRYPT_MODE, secKey, new IvParameterSpec(password.getBytes()));
            byte[] encBytes = cipher.doFinal(clearText.getBytes());
            return new String(Base64.getEncoder().encode(encBytes));
        } catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException | IllegalBlockSizeException | BadPaddingException | InvalidAlgorithmParameterException e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * AES密码分组链接模式Cipher Block Chaining 解密
     *
     * @param password   输入的对称加密秘钥
     * @param cipherText 需要解密的内容
     * @return 解密之后的明文
     */
    public static String aesCBCDecrypt(String password, String cipherText) {
        password = getPre16OfValue(password);
        try {
            Cipher cipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
            SecretKey secKey = new SecretKeySpec(password.getBytes(), "AES");
            cipher.init(Cipher.DECRYPT_MODE, secKey, new IvParameterSpec(password.getBytes()));
            byte[] encBytes = Base64.getDecoder().decode(cipherText.getBytes());
            byte[] decBytes = cipher.doFinal(encBytes);
            return new String(decBytes);
        } catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException | IllegalBlockSizeException | BadPaddingException | InvalidAlgorithmParameterException e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * AES 电码本加密 Electronic Codebook Book
     *
     * @param password  秘钥
     * @param clearText 明文
     * @return 加密的密文
     */
    public static String aesECBEncrypt(String password, String clearText) {
        password = getPre16OfValue(password);
        try {
            Cipher cipher = Cipher.getInstance("AES/ECB/PKCS5Padding");
            SecretKey secKey = new SecretKeySpec(password.getBytes(), "AES");
            cipher.init(Cipher.ENCRYPT_MODE, secKey);
            byte[] encBytes = cipher.doFinal(clearText.getBytes());
            return new String(Base64.getEncoder().encode(encBytes));
        } catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException | IllegalBlockSizeException | BadPaddingException e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * AES 电码本加密 Electronic Codebook Book 解密方法
     *
     * @param password   密码
     * @param cipherText 密文
     * @return 解密后的内容
     */
    public static String aesECBDecrypt(String password, String cipherText) {
        password = getPre16OfValue(password);
        try {
            Cipher cipher = Cipher.getInstance("AES/ECB/PKCS5Padding");
            SecretKey secKey = new SecretKeySpec(password.getBytes(), "AES");
            cipher.init(Cipher.DECRYPT_MODE, secKey);
            byte[] encBytes = Base64.getDecoder().decode(cipherText.getBytes());
            byte[] decBytes = cipher.doFinal(encBytes);
            return new String(decBytes);
        } catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException | IllegalBlockSizeException | BadPaddingException e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * 这部分是给内存RSA使用的，不推荐跨机器使用
     * 正确做法应该是使用java keytool 生成密钥文件jks，然后私钥持有端通过读取jks文件，加解密
     * 公钥持有端，通过颁发的证书（通常是.cert）文件进行加解密
     */
    public static class RSAKeyPair {
        private RSAPublicKey rsaPublicKey;
        private RSAPrivateKey rsaPrivateKey;

        public RSAPrivateKey getRsaPrivateKey() {
            return rsaPrivateKey;
        }

        public RSAPublicKey getRsaPublicKey() {
            return rsaPublicKey;
        }
    }

    /**
     * RSA非对称加密，使用指定的密码获取公私钥
     *
     * @param password 密码明文
     * @return 一个公私钥的对象，可以从此对象获取公私钥
     */
    public static RSAKeyPair newRSAKeyPair(String password) {
        try {
            KeyPairGenerator gen = KeyPairGenerator.getInstance("RSA");
            SecureRandom random = new SecureRandom(password.getBytes());
            gen.initialize(1024, random);
            KeyPair store = gen.generateKeyPair();
            RSAKeyPair rsaKeyPair = new RSAKeyPair();
            rsaKeyPair.rsaPrivateKey = (RSAPrivateKey) store.getPrivate();
            rsaKeyPair.rsaPublicKey = (RSAPublicKey) store.getPublic();
            return rsaKeyPair;
        } catch (NoSuchAlgorithmException e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * 非对称加密使用私钥加密
     *
     * @param rsaPrivateKey 获取公私钥对儿之后的私钥字段
     * @param clearText     明文
     * @return 加密文本
     */
    public static String rsaPrivateKeyEncode(RSAPrivateKey rsaPrivateKey, String clearText) {
        byte[] clearBytes = clearText.getBytes();
        if (clearBytes.length > 117) {
            throw new RuntimeException("Input String too long, the max size is 117 bytes.");
        }
        try {
            PKCS8EncodedKeySpec pkcs8EncodedKeySpec = new PKCS8EncodedKeySpec(rsaPrivateKey.getEncoded());
            KeyFactory keyFactory = KeyFactory.getInstance("RSA");
            PrivateKey privateKey = keyFactory.generatePrivate(pkcs8EncodedKeySpec);
            Cipher cipher = Cipher.getInstance("RSA");
            cipher.init(Cipher.ENCRYPT_MODE, privateKey);
            byte[] encBytes = cipher.doFinal(clearBytes);
            return new String(Base64.getEncoder().encode(encBytes));
        } catch (NoSuchAlgorithmException | InvalidKeySpecException | NoSuchPaddingException | InvalidKeyException | IllegalBlockSizeException | BadPaddingException e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * 非对称加密使用公钥解密
     *
     * @param rsaPublicKey 公钥
     * @param cipherText   加密后的字符串
     * @return 明文
     */
    public static String rsaPublicKeyDecode(RSAPublicKey rsaPublicKey, String cipherText) {
        try {
            X509EncodedKeySpec x509EncodedKeySpec = new X509EncodedKeySpec(rsaPublicKey.getEncoded());
            KeyFactory keyFactory = KeyFactory.getInstance("RSA");
            PublicKey publicKey = keyFactory.generatePublic(x509EncodedKeySpec);
            Cipher cipher = Cipher.getInstance("RSA");
            cipher.init(Cipher.DECRYPT_MODE, publicKey);
            byte[] encBytes = Base64.getDecoder().decode(cipherText.getBytes());
            byte[] clearBytes = cipher.doFinal(encBytes);
            return new String(clearBytes);
        } catch (InvalidKeySpecException | NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException | IllegalBlockSizeException | BadPaddingException e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * 非对称加密使用公钥加密
     *
     * @param rsaPublicKey 公钥
     * @param clearText    明文
     * @return 密文
     */
    public static String rsaPublicKeyEncode(RSAPublicKey rsaPublicKey, String clearText) {
        byte[] clearBytes = clearText.getBytes();
        if (clearBytes.length > 117) {
            throw new RuntimeException("Input String too long, the max size is 117 bytes.");
        }
        try {
            X509EncodedKeySpec x509EncodedKeySpec = new X509EncodedKeySpec(rsaPublicKey.getEncoded());
            KeyFactory keyFactory = KeyFactory.getInstance("RSA");
            PublicKey publicKey = keyFactory.generatePublic(x509EncodedKeySpec);
            Cipher cipher = Cipher.getInstance("RSA");
            cipher.init(Cipher.ENCRYPT_MODE, publicKey);
            byte[] encBytes = cipher.doFinal(clearBytes);
            return new String(Base64.getEncoder().encode(encBytes));
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }
    }

    /**
     * 私钥解密
     *
     * @param rsaPrivateKey 私钥对象
     * @param cipherText    密文
     * @return 明文
     */
    public static String rsaPrivateKeyDecode(RSAPrivateKey rsaPrivateKey, String cipherText) {
        try {
            PKCS8EncodedKeySpec pkcs8EncodedKeySpec = new PKCS8EncodedKeySpec(rsaPrivateKey.getEncoded());
            KeyFactory keyFactory = KeyFactory.getInstance("RSA");
            PrivateKey privateKey = keyFactory.generatePrivate(pkcs8EncodedKeySpec);
            Cipher cipher = Cipher.getInstance("RSA");
            cipher.init(Cipher.DECRYPT_MODE, privateKey);
            byte[] encBytes = Base64.getDecoder().decode(cipherText.getBytes());
            byte[] clearText = cipher.doFinal(encBytes);
            return new String(clearText);
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }
    }

    /**
     * 从本地JKS文件中获取私钥的方法
     *
     * @param jksFilePath   jks文件路径
     * @param alias         密钥对的别名
     * @param jksPassword   仓库密码
     * @param aliasPassword 密钥对密码
     * @return 私钥对象
     */
    public static RSAPrivateKey fetchRSAPrivateKeyFromJKSFile(String jksFilePath, String alias, String jksPassword, String aliasPassword) {
        try (InputStream is = new FileInputStream(jksFilePath)) {
            KeyStore store = KeyStore.getInstance("JKS");
            store.load(is, jksPassword.toCharArray());
            RSAPrivateKey privateKey = (RSAPrivateKey) store.getKey(alias, aliasPassword.toCharArray());
            return privateKey;
        } catch (KeyStoreException | IOException | NoSuchAlgorithmException | CertificateException | UnrecoverableKeyException e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * 从证书文件中获取公钥，生成public key对象
     * @param certFilePath 证书的路径
     * @return 公钥对象
     */
    public static RSAPublicKey fetchRSAPublicKeyFromCertFile(String certFilePath) {
        try (InputStream is = new FileInputStream(certFilePath)) {
            CertificateFactory cf = CertificateFactory.getInstance("X.509");
            X509Certificate cert = (X509Certificate) cf.generateCertificate(is);
            RSAPublicKey publicKey = (RSAPublicKey) cert.getPublicKey();
            return publicKey;
        } catch (CertificateException | IOException e) {
            e.printStackTrace();
        }
        return null;
    }
}
